package com.zkl.evaluation.manager;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.*;
import com.zkl.evaluation.common.ErrorCode;
import com.zkl.evaluation.config.AiConfig;
import com.zkl.evaluation.exception.BusinessException;
import io.reactivex.Flowable;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

@Component
@Slf4j
public class AiManager {
    @Resource
    private OkHttpClient client;

    @Resource
    private AiConfig aiConfig;

    @Resource
    private ClientV4 clientV4;

    private List<RoleContent> historyList = new ArrayList<>(); // 对话历史
    private String totalAnswer = ""; // 大模型的答案汇总
    private boolean wsCloseFlag = false; // WebSocket 关闭标志
    private CountDownLatch latch; // 同步锁

    /**
     * 通用请求方法（支持 systemMessage 和 userMessage）
     *
     * @param systemMessage 系统消息（用于设置对话背景或指令）
     * @param userMessage   用户消息（用户输入的问题或内容）
     * @return 大模型的响应结果
     */
    public String doRequest(String systemMessage, String userMessage) throws Exception {
        this.historyList = new ArrayList<>();
        this.totalAnswer = ""; // 清空之前的答案
        this.wsCloseFlag = false; // 重置关闭标志
        this.latch = new CountDownLatch(1); // 初始化同步锁

        // 构建鉴权 URL
        String authUrl = aiConfig.authUrl();
        String url = authUrl.replace("http://", "ws://").replace("https://", "wss://");

        // 创建 WebSocket 请求
        Request request = new Request.Builder().url(url).build();
        WebSocket webSocket = client.newWebSocket(request, new WebSocketListener() {
            @Override
            public void onOpen(WebSocket webSocket, Response response) {
                // 发送消息
                sendMessage(webSocket, systemMessage, userMessage);
            }

            @Override
            public void onMessage(WebSocket webSocket, String text) {
                handleResponse(text);
            }

            @Override
            public void onFailure(WebSocket webSocket, Throwable t, Response response) {
                t.printStackTrace();
                latch.countDown(); // 异常时释放锁
            }
        });

        // 等待响应完成（设置超时时间）
        if (!latch.await(30, TimeUnit.SECONDS)) {
            throw new RuntimeException("Request timed out");
        }

        return totalAnswer;
    }

    /**
     * 发送消息到 WebSocket（支持 systemMessage 和 userMessage）
     */
    public static final String Spark4_Ultra = Constants.ModelChatGLM4;
    private void sendMessage(WebSocket webSocket, String systemMessage, String userMessage) {
        try {
            JSONObject requestJson = new JSONObject();

            // 构建 header
            JSONObject header = new JSONObject();
            header.put("app_id", aiConfig.getAppid());
            header.put("uid", UUID.randomUUID().toString().substring(0, 10));

            // 构建 parameter
            JSONObject parameter = new JSONObject();
            JSONObject chat = new JSONObject();
            chat.put("domain", aiConfig.getDomain());
            chat.put("temperature", 0.5); // 默认温度值
            chat.put("max_tokens", 4096); // 最大 token 数
            parameter.put("chat", chat);

            // 构建 payload
            JSONObject payload = new JSONObject();
            JSONObject messageJson = new JSONObject();
            JSONArray text = new JSONArray();

            // 添加系统消息
            if (systemMessage != null && !systemMessage.isEmpty()) {
                RoleContent systemRoleContent = new RoleContent();
                systemRoleContent.setRole("system");
                systemRoleContent.setContent(systemMessage);
                text.add(JSON.toJSON(systemRoleContent));
                historyList.add(systemRoleContent);
            }

            // 添加用户消息
            RoleContent userRoleContent = new RoleContent();
            userRoleContent.setRole("user");
            userRoleContent.setContent(userMessage);
            text.add(JSON.toJSON(userRoleContent));
            historyList.add(userRoleContent);

            messageJson.put("text", text);
            payload.put("message", messageJson);

            // 组装请求
            requestJson.put("header", header);
            requestJson.put("parameter", parameter);
            requestJson.put("payload", payload);

            // 发送请求
            webSocket.send(requestJson.toString());
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 处理 WebSocket 响应
     */
    private void handleResponse(String text) {
        try {
            JSONObject responseJson = JSON.parseObject(text);
            JSONObject header = responseJson.getJSONObject("header");
            if (header.getInteger("code") != 0) {
                System.err.println("Error: " + header.getString("message"));
                return;
            }

            // 解析响应内容
            JSONObject payload = responseJson.getJSONObject("payload");
            JSONObject choices = payload.getJSONObject("choices");
            JSONArray textArray = choices.getJSONArray("text");
            for (int i = 0; i < textArray.size(); i++) {
                JSONObject textObj = textArray.getJSONObject(i);
                totalAnswer += textObj.getString("content");
            }

            // 如果状态为 2，表示响应完成
            if (header.getInteger("status") == 2) {
                // 将大模型的回答加入历史
                RoleContent assistantResponse = new RoleContent();
                assistantResponse.setRole("assistant");
                assistantResponse.setContent(totalAnswer);
                historyList.add(assistantResponse);

                wsCloseFlag = true;
                latch.countDown(); // 释放锁
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }



    // 稳定的随机数
    private static final float STABLE_TEMPERATURE = 0.05f;

    // 不稳定的随机数
    private static final float UNSTABLE_TEMPERATURE = 0.99f;

    /**
     * 同步请求（答案不稳定）
     *
     * @param systemMessage
     * @param userMessage
     * @return
     */
    public String doSyncUnstableRequest(String systemMessage, String userMessage) {
        return doRequest(systemMessage, userMessage, Boolean.FALSE, UNSTABLE_TEMPERATURE);
    }

    /**
     * 同步请求（答案较稳定）
     *
     * @param systemMessage
     * @param userMessage
     * @return
     */
    public String doSyncStableRequest(String systemMessage, String userMessage) {
        log.info(systemMessage);
        log.info(userMessage);
        return doRequest(systemMessage, userMessage, Boolean.FALSE, STABLE_TEMPERATURE);
    }

    /**
     * 同步请求
     *
     * @param systemMessage
     * @param userMessage
     * @param temperature
     * @return
     */
    public String doSyncRequest(String systemMessage, String userMessage, Float temperature) {
        return doRequest(systemMessage, userMessage, Boolean.FALSE, temperature);
    }

    /**
     * 通用请求（简化消息传递）
     *
     * @param systemMessage
     * @param userMessage
     * @param stream
     * @param temperature
     * @return
     */
    public String doRequest(String systemMessage, String userMessage, Boolean stream, Float temperature) {
        List<ChatMessage> chatMessageList = new ArrayList<>();
        ChatMessage systemChatMessage = new ChatMessage(ChatMessageRole.SYSTEM.value(), systemMessage);
        chatMessageList.add(systemChatMessage);
        ChatMessage userChatMessage = new ChatMessage(ChatMessageRole.USER.value(), userMessage);
        chatMessageList.add(userChatMessage);
        return doRequest(chatMessageList, stream, temperature);
    }

    /**
     * 通用请求
     *
     * @param messages
     * @param stream
     * @param temperature
     * @return
     */
    public String doRequest(List<ChatMessage> messages, Boolean stream, Float temperature) {
        // 构建请求
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Spark4_Ultra)
                .stream(stream)
                .temperature(temperature)
                .invokeMethod(Constants.invokeMethod)
                .messages(messages)
                .build();
        try {
            ModelApiResponse invokeModelApiResp = clientV4.invokeModelApi(chatCompletionRequest);
            return invokeModelApiResp.getData().getChoices().get(0).toString();
        } catch (Exception e) {
            e.printStackTrace();
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, e.getMessage());
        }
    }

    /**
     * 通用流式请求（简化消息传递）
     *
     * @param systemMessage
     * @param userMessage
     * @param temperature
     * @return
     */
    public Flowable<ModelData> doStreamRequest(String systemMessage, String userMessage, Float temperature) {
        List<ChatMessage> chatMessageList = new ArrayList<>();
        ChatMessage systemChatMessage = new ChatMessage(ChatMessageRole.SYSTEM.value(), systemMessage);
        chatMessageList.add(systemChatMessage);
        ChatMessage userChatMessage = new ChatMessage(ChatMessageRole.USER.value(), userMessage);
        chatMessageList.add(userChatMessage);
        return doStreamRequest(chatMessageList, temperature);
    }


    /**
     * 通用流式请求
     *
     * @param messages
     * @param temperature
     * @return
     */
    public Flowable<ModelData> doStreamRequest(List<ChatMessage> messages, Float temperature) {
        // 构建请求
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Spark4_Ultra)
                .stream(Boolean.TRUE)
                .temperature(temperature)
                .invokeMethod(Constants.invokeMethod)
                .messages(messages)
                .build();
        try {
            ModelApiResponse invokeModelApiResp = clientV4.invokeModelApi(chatCompletionRequest);
            return invokeModelApiResp.getFlowable();
        } catch (Exception e) {
            e.printStackTrace();
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, e.getMessage());
        }
    }

    /**
     * 内部类：角色内容
     */
    private static class RoleContent {
        private String role;
        private String content;

        public String getRole() {
            return role;
        }

        public void setRole(String role) {
            this.role = role;
        }

        public String getContent() {
            return content;
        }

        public void setContent(String content) {
            this.content = content;
        }
    }
}